{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🔥 LightlyTrain - Semantic Segmentation with DINOv3 EoMT 🔥\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for semantic segmentation with our state-of-the-art [EoMT](https://arxiv.org/abs/2503.19108) model built on [DINOv3](https://github.com/facebookresearch/dinov3) backbones, with our publicly released weights trained on the [COCO-Stuff](https://arxiv.org/abs/1612.03716) and [Cityscapes](https://www.cityscapes-dataset.com/) dataset.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/eomt_semantic_segmentation.ipynb)\n",
    "\n",
    "> **Important**: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to `Runtime` > `Change runtime type` and selecting a GPU hardware accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "LightlyTrain can be installed directly via `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Important**: LightlyTrain is officially supported on\n",
    "> - Linux: CPU or CUDA\n",
    "> - MacOS: CPU only\n",
    "> - Windows (experimental): CPU or CUDA\n",
    ">\n",
    "> We are planning to support MPS for MacOS.\n",
    ">\n",
    "> Check the [installation instructions](https://docs.lightly.ai/train/stable/installation.html) for more details on installation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction using LightlyTrain's model weights\n",
    "\n",
    "### Download an example image\n",
    "\n",
    "Download an example image for inference with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O cat.jpg https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the model weights to LightlyTrain\n",
    "\n",
    "Then, we load the model weights with LightlyTrain's `load_model` function and do inference on the example image with the `predict` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "model = lightly_train.load_model(\"dinov3/vits16-eomt-coco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict the mask\n",
    "\n",
    "Simply call the `model`'s `.predict()` method to predict the mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masks = model.predict(\"cat.jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the image and mask\n",
    "\n",
    "Finally, we visualize the image and mask to see if it makes sense:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torchvision.io import read_image\n",
    "from torchvision.utils import draw_segmentation_masks\n",
    "\n",
    "image = read_image(\"cat.jpg\")\n",
    "masks = torch.stack([masks == class_id for class_id in masks.unique()])\n",
    "image_with_masks = draw_segmentation_masks(image, masks, alpha=0.6)\n",
    "plt.imshow(image_with_masks.permute(1, 2, 0))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congrats! Using LightlyTrain's EoMT model weights for inference is as simple as the above example."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
